import os, pickle, logging, numpy as np
from tqdm import tqdm
import math

class NTU_Generator():
    def __init__(self, args):
        self.num_person_out = 2
        self.num_person_in = 4
        self.num_joint = 25
        self.max_frame = 300
        self.print_bar = True

        self.dataset = args.dataset
        self.data_path = args.data_path
        assert type(self.data_path)==list,'data_path shoule be a list type,please check it'
        self.pre_normalization = args.pre_normalization

        self.ignored_sample_path = args.ignored_path
        # ntu_ignored = '{}/ignore.txt'.format(self.ignored_sample_path)
        ntu_ignored = os.path.join(args.ignored_path,'ignore.txt')
        if os.path.exists(ntu_ignored):
            with open(ntu_ignored, 'r') as f:
                self.ignored_samples = [
                    line.strip() + '.skeleton' for line in f.readlines()
                ]
        else:
            logging.info('')
            logging.error('Error: Do NOT exist ignored sample file {}'.format(ntu_ignored))
            self.ignored_samples = []

        self.out_path = '{}/{}'.format(args.out_path, self.dataset.replace('-', '/'))
        if not os.path.exists(self.out_path):
            os.makedirs(self.out_path)

        # Divide train and eval samples
        training_samples = dict()
        training_samples['ntu-xsub'] = [
            1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38
        ]
        training_samples['ntu-xview'] = [2, 3]
        training_samples['ntu-xsub120'] = [
            1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35,
            38, 45, 46, 47, 49, 50, 52, 53, 54, 55, 56, 57, 58, 59, 70, 74, 78,
            80, 81, 82, 83, 84, 85, 86, 89, 91, 92, 93, 94, 95, 97, 98, 100, 103
        ]
        training_samples['ntu-xset120'] = set(range(2, 33, 2))
        self.training_sample = training_samples[self.dataset]
        # Get skeleton file list
        self.file_list = []

        for folder in self.data_path:
            for filename in os.listdir(folder):
                self.file_list.append((folder, filename))
            # if '120' not in self.dataset:  # for NTU 60, only one folder
            #     break

    def start(self):
        # Generate data
        for phase in ['train', 'eval']:
            logging.info('Phase: {}'.format(phase))
            self.gendata(phase)

    def gendata(self, phase):
        sample_name = []
        sample_label = []
        sample_paths = []
        for folder, filename in sorted(self.file_list):
            if filename in self.ignored_samples:
                continue
            path = os.path.join(folder, filename)
            setup_id = int(filename[(filename.find('S') + 1):(filename.find('S') + 4)])
            action_class = int(filename[(filename.find('A') + 1):(filename.find('A') + 4)])
            subject_id = int(filename[(filename.find('P') + 1):(filename.find('P') + 4)])
            camera_id = int(filename[(filename.find('C') + 1):(filename.find('C') + 4)])

            if self.dataset == 'ntu-xview':
                istraining = (camera_id in self.training_sample)
            elif self.dataset == 'ntu-xsub' or self.dataset == 'ntu-xsub120':
                istraining = (subject_id in self.training_sample)
            elif self.dataset == 'ntu-xset120':
                istraining = (setup_id in self.training_sample)
            else:
                istraining = None
                logging.info('')
                logging.error('Error: Do NOT exist this dataset {}'.format(self.dataset))
                raise ValueError()

            if phase == 'train':
                issample = istraining
            elif phase == 'eval':
                issample = not (istraining)
            else:
                issample = None
                logging.info('')
                logging.error('Error: Do NOT exist this phase {}'.format(phase))
                raise ValueError()

            if issample:
                sample_paths.append(path)
                sample_label.append(action_class - 1)  # to 0-indexed

        # Save labels
        with open('{}/{}_label.pkl'.format(self.out_path, phase), 'wb') as f:
            pickle.dump((sample_paths, list(sample_label)), f)

        if self.pre_normalization:
            # Create data tensor (N,C,T,V,M)
            fp = np.zeros((len(sample_label), 3, self.max_frame, self.num_joint, self.num_person_out), dtype=np.float32)
            # Fill (C,T,V,M) to data tensor (N,C,T,V,M)
            items = tqdm(sample_paths, dynamic_ncols=True) if self.print_bar else sample_paths
            for i, s in enumerate(items):
                data = self.read_xyz(s)
                fp[i, :, 0:data.shape[1], :, :] = data
            # Perform preprocessing on data tensor
            fp = pre_normalization(fp, print_bar=self.print_bar)
            # Save input data (train/eval)
            np.save('{}/{}_data.npy'.format(self.out_path, phase), fp)
        else:
            fp = np.zeros((len(sample_label), 3, self.max_frame, self.num_joint, self.num_person_out), dtype=np.float32)
            items = tqdm(sample_paths, dynamic_ncols=True) if self.print_bar else sample_paths
            for i, s in enumerate(items):
                data = self.read_org_xyz(s)
                fp[i, :, 0:data.shape[1], :, :] = data
            np.save('{}/{}_data.npy'.format(self.out_path, phase), fp)


    def read_skeleton_filter(self, file):
        with open(file, 'r') as f:
            skeleton_sequence = {}
            skeleton_sequence['numFrame'] = int(f.readline())
            skeleton_sequence['frameInfo'] = []
            for t in range(skeleton_sequence['numFrame']):
                frame_info = {}
                frame_info['numBody'] = int(f.readline())
                frame_info['bodyInfo'] = []

                for m in range(frame_info['numBody']):
                    body_info = {}
                    body_info_key = [
                        'bodyID', 'clipedEdges', 'handLeftConfidence',
                        'handLeftState', 'handRightConfidence', 'handRightState',
                        'isResticted', 'leanX', 'leanY', 'trackingState'
                    ]
                    body_info = {
                        k: float(v)
                        for k, v in zip(body_info_key, f.readline().split())
                    }
                    body_info['numJoint'] = int(f.readline())
                    body_info['jointInfo'] = []
                    for v in range(body_info['numJoint']):
                        joint_info_key = [
                            'x', 'y', 'z', 'depthX', 'depthY', 'colorX', 'colorY',
                            'orientationW', 'orientationX', 'orientationY',
                            'orientationZ', 'trackingState'
                        ]
                        joint_info = {
                            k: float(v)
                            for k, v in zip(joint_info_key, f.readline().split())
                        }
                        body_info['jointInfo'].append(joint_info)
                    frame_info['bodyInfo'].append(body_info)
                skeleton_sequence['frameInfo'].append(frame_info)

        return skeleton_sequence


    def get_nonzero_std(self, s):  # (T,V,C)
        index = s.sum(-1).sum(-1) != 0  # select valid frames
        s = s[index]
        if len(s) != 0:
            s = s[:, :, 0].std() + s[:, :, 1].std() + s[:, :, 2].std()  # three channels
        else:
            s = 0
        return s


    def read_xyz(self, file):
        seq_info = self.read_skeleton_filter(file)
        data = np.zeros((self.num_person_in, seq_info['numFrame'], self.num_joint, 3))
        for n, f in enumerate(seq_info['frameInfo']):
            for m, b in enumerate(f['bodyInfo']):
                for j, v in enumerate(b['jointInfo']):
                    if m < self.num_person_in and j < self.num_joint:
                        data[m, n, j, :] = [v['x'], v['y'], v['z']]

        # select two max energy body
        energy = np.array([self.get_nonzero_std(x) for x in data])
        index = energy.argsort()[::-1][0:self.num_person_out]
        data = data[index]

        data = data.transpose(3, 1, 2, 0)  # to (C,T,V,M)
        return data

    def read_org_xyz(self, file):

        data = np.zeros((3, self.max_frame, self.num_joint, self.num_person_out))
        with open(file, 'r') as fr:
            frame_num = int(fr.readline())
            for frame in range(frame_num):
                if frame >= self.max_frame:
                    break
                person_num = int(fr.readline())
                for person in range(person_num):
                    fr.readline()
                    joint_num = int(fr.readline())
                    for joint in range(joint_num):
                        v = fr.readline().split(' ')
                        if joint < self.num_joint and person < self.num_person_out:
                            data[0, frame, joint, person] = float(v[0])
                            data[1, frame, joint, person] = float(v[1])
                            data[2, frame, joint, person] = float(v[2])
        return data



# Thanks to SHI Lei for the released code on Github (https://github.com/lshiwjx/2s-AGCN)
def pre_normalization(data, print_bar=True, zaxis=[0, 1], xaxis=[8, 4]):
    N, C, T, V, M = data.shape
    s = np.transpose(data, [0, 4, 2, 3, 1])  # N, C, T, V, M  to  N, M, T, V, C

    logging.info('Pad the null frames with the previous frames')
    items = tqdm(s, dynamic_ncols=True) if print_bar else s
    for i_s, skeleton in enumerate(items):  # pad
        if skeleton.sum() == 0:
            logging.info('Sample {:d} has no skeleton'.format(i_s))
        for i_p, person in enumerate(skeleton):
            if person.sum() == 0:
                continue
            if person[0].sum() == 0:
                index = (person.sum(-1).sum(-1) != 0)
                tmp = person[index].copy()
                person *= 0
                person[:len(tmp)] = tmp
            for i_f, frame in enumerate(person):
                if frame.sum() == 0:
                    if person[i_f:].sum() == 0:
                        rest = len(person) - i_f
                        num = int(np.ceil(rest / i_f))
                        pad = np.concatenate([person[0:i_f] for _ in range(num)], 0)[:rest]
                        s[i_s, i_p, i_f:] = pad
                        break

    logging.info('Sub the center joint #1 (spine joint in ntu and neck joint in kinetics)')
    items = tqdm(s, dynamic_ncols=True) if print_bar else s
    for i_s, skeleton in enumerate(items):
        if skeleton.sum() == 0:
            continue
        main_body_center = skeleton[0][:, 1:2, :].copy()
        for i_p, person in enumerate(skeleton):
            if person.sum() == 0:
                continue
            mask = (person.sum(-1) != 0).reshape(T, V, 1)
            s[i_s, i_p] = (s[i_s, i_p] - main_body_center) * mask

    logging.info('Parallel the bone between hip(jpt 0) and spine(jpt 1) of the first person to the z axis')
    items = tqdm(s, dynamic_ncols=True) if print_bar else s
    for i_s, skeleton in enumerate(items):
        if skeleton.sum() == 0:
            continue
        joint_bottom = skeleton[0, 0, zaxis[0]]
        joint_top = skeleton[0, 0, zaxis[1]]
        axis = np.cross(joint_top - joint_bottom, [0, 0, 1])
        angle = angle_between(joint_top - joint_bottom, [0, 0, 1])
        matrix_z = rotation_matrix(axis, angle)
        for i_p, person in enumerate(skeleton):
            if person.sum() == 0:
                continue
            for i_f, frame in enumerate(person):
                if frame.sum() == 0:
                    continue
                for i_j, joint in enumerate(frame):
                    s[i_s, i_p, i_f, i_j] = np.dot(matrix_z, joint)

    logging.info('Parallel the bone between right shoulder(jpt 8) and left shoulder(jpt 4) of the first person to the x axis')
    items = tqdm(s, dynamic_ncols=True) if print_bar else s
    for i_s, skeleton in enumerate(items):
        if skeleton.sum() == 0:
            continue
        joint_rshoulder = skeleton[0, 0, xaxis[0]]
        joint_lshoulder = skeleton[0, 0, xaxis[1]]
        axis = np.cross(joint_rshoulder - joint_lshoulder, [1, 0, 0])
        angle = angle_between(joint_rshoulder - joint_lshoulder, [1, 0, 0])
        matrix_x = rotation_matrix(axis, angle)
        for i_p, person in enumerate(skeleton):
            if person.sum() == 0:
                continue
            for i_f, frame in enumerate(person):
                if frame.sum() == 0:
                    continue
                for i_j, joint in enumerate(frame):
                    s[i_s, i_p, i_f, i_j] = np.dot(matrix_x, joint)

    data = np.transpose(s, [0, 4, 2, 3, 1])
    return data


def rotation_matrix(axis, theta):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.
    """
    if np.abs(axis).sum() < 1e-6 or np.abs(theta) < 1e-6:
        return np.eye(3)
    axis = np.asarray(axis)
    axis = axis / math.sqrt(np.dot(axis, axis))
    a = math.cos(theta / 2.0)
    b, c, d = -axis * math.sin(theta / 2.0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
                     [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
                     [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])


def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)


def angle_between(v1, v2):
    """ Returns the angle in radians between vectors 'v1' and 'v2'::

            >>> angle_between((1, 0, 0), (0, 1, 0))
            1.5707963267948966
            >>> angle_between((1, 0, 0), (1, 0, 0))
            0.0
            >>> angle_between((1, 0, 0), (-1, 0, 0))
            3.141592653589793
    """
    if np.abs(v1).sum() < 1e-6 or np.abs(v2).sum() < 1e-6:
        return 0
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))


def x_rotation(vector, theta):
    """Rotates 3-D vector around x-axis"""
    R = np.array([[1, 0, 0], [0, np.cos(theta), -np.sin(theta)], [0, np.sin(theta), np.cos(theta)]])
    return np.dot(R, vector)


def y_rotation(vector, theta):
    """Rotates 3-D vector around y-axis"""
    R = np.array([[np.cos(theta), 0, np.sin(theta)], [0, 1, 0], [-np.sin(theta), 0, np.cos(theta)]])
    return np.dot(R, vector)


def z_rotation(vector, theta):
    """Rotates 3-D vector around z-axis"""
    R = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
    return np.dot(R, vector)
